import cv2
import torch.cuda
from super_gradients.training import models
from super_gradients.common.object_names import Models


# model = models.get('yolo_nas_s', num_classes=1, checkpoint_path="")


# model = model.to('cuda' if torch.cuda.is_available() else 'cpu')
# 打开摄像头并检测
# model.predict_webcam()

# 检测图片
# img = cv2.imread("2.jpg")
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# outputs = model.predict(img)
# outputs.show()


def to_onnx():
    model = models.get(Models.YOLO_NAS_S, pretrained_weights='coco')
    # 导出,需要使cpu方式
    models.convert_to_onnx(model=model, input_shape=(3, 640, 640), out_path="yolo_nas_s.onnx")


def predict_webcam():
    model = models.get(Models.YOLO_NAS_S, pretrained_weights='coco')
    model = model.to('cuda' if torch.cuda.is_available() else 'cpu')
    model.predict_webcam()


def predict_image():
    model = models.get(Models.YOLO_NAS_S, checkpoint_path=r"test\1400\ckpt_best.pth", pretrained_weights='coco'
                       , num_classes=1)
    model = model.to('cuda' if torch.cuda.is_available() else 'cpu')
    outputs = model.predict("t1.jpg", iou=0.6, conf=0.3, fuse_model=False)
    outputs.show()


if __name__ == "__main__":
    predict_image()
